﻿using System;
using System.Disposables;
using System.Linq;

namespace RxPowerToys.ObservableExtensions
{
    /// <summary>
    /// Contains the following IObservable<T> extension methods: 
    /// WhenFirst
    /// WhenFirstOrDefault
    /// WhenLast
    /// WhenLastOrDefault
    /// WhenCount
    /// WhenIsEmpty
    /// WhenAggregated
    /// WhenSingle
    /// WhenAll
    /// WhenAny
    /// </summary>
    public static class AsynchronousAggregation
    {
        public static IObservable<T> WhenFirst<T>(this IObservable<T> source)
        {
            return WhenFirst(source, true);
        }

        public static IObservable<T> WhenFirstOrDefault<T>(this IObservable<T> source)
        {
            return WhenFirst(source, false);
        }

        private static IObservable<T> WhenFirst<T>(this IObservable<T> source, bool errorOnEmpty)
        {
            return Observable.CreateWithDisposable<T>(o =>
            {
                bool valueReceived = false;
                var subscription = new FutureDisposable();
                subscription.Set(source.Subscribe(n =>
                {
                    if (!valueReceived)
                    {
                        valueReceived = true;
                        o.OnNext(n);
                        o.OnCompleted();
                    }
                    subscription.Dispose(); // we have what we want, avoid wasting resources
                },
                    e =>
                    {
                        if (!valueReceived)
                            o.OnError(e);
                    },
                    () =>
                    {
                        if (!valueReceived)
                        {
                            if (errorOnEmpty)
                                o.OnError(new InvalidOperationException("Sequence contains no elements"));
                            else
                                o.OnNext(default(T));
                        }
                    }
                ));
                return subscription;
            });
        }

        public static IObservable<T> WhenLast<T>(this IObservable<T> source)
        {
            return WhenLast(source, true);
        }

        public static IObservable<T> WhenLastOrDefault<T>(this IObservable<T> source)
        {
            return WhenLast(source, false);
        }

        private static IObservable<T> WhenLast<T>(this IObservable<T> source, bool errorOnEmpty)
        {
            return Observable.CreateWithDisposable<T>(o =>
            {
                bool valueReceived = false;
                T lastValue = default(T);
                return source.Subscribe(n =>
                {
                    valueReceived = true;
                    lastValue = n;
                },
                o.OnError,
                () =>
                {
                    if (!valueReceived && errorOnEmpty)
                        o.OnError(new InvalidOperationException("Sequence contains no elements"));
                    else
                    {
                        o.OnNext(valueReceived ? lastValue : default(T));
                        o.OnCompleted();
                    }
                });
            });
        }

        public static IObservable<int> WhenCount<T>(this IObservable<T> source)
        {
            return source.Scan<T, int>(0, (count, n) => count + 1).WhenLast();
        }

        public static IObservable<Unit> WhenIsEmpty<T>(this IObservable<T> source)
        {
            return Observable.CreateWithDisposable<Unit>(o =>
            {
                bool valueReceived = false;
                var subscription = new FutureDisposable();
                subscription.Set(source.Subscribe(n =>
                {
                    if (!valueReceived)
                    {
                        valueReceived = true;
                        o.OnCompleted();
                    }
                    subscription.Dispose(); // we have what we want, avoid wasting resources
                },
                o.OnError,
                () =>
                {
                    if (!valueReceived)
                    {
                        o.OnNext(new Unit());
                        o.OnCompleted();
                    }
                }));
                return subscription;
            });
        }

        public static IObservable<T> WhenAggregated<T>(this IObservable<T> source, Func<T, T, T> accumulator)
        {
            return Observable.Scan(source, accumulator).WhenLast();
        }

        public static IObservable<TAccumulate> WhenAggregated<T, TAccumulate>(this IObservable<T> source, TAccumulate seed, Func<TAccumulate, T, TAccumulate> accumulator)
        {
            return Observable.Scan(source, seed, accumulator).WhenLast();
        }

        private static IObservable<T> WhenSingle<T>(this IObservable<T> source, bool errorOnMoreOrEmpty)
        {
            return Observable.CreateWithDisposable<T>(o =>
            {
                int count = 0;
                T firstValue = default(T);
                var subscription = new FutureDisposable();
                subscription.Set(source.Subscribe(n =>
                {
                    if (count == 0)
                        firstValue = n;
                    ++count;
                    if (count == 2)
                    {
                        if (errorOnMoreOrEmpty)
                            o.OnError(new InvalidOperationException("Sequence contains more than one element"));
                        else
                        {
                            o.OnNext(default(T));
                            o.OnCompleted();
                        }
                    }

                    if (count > 1)
                        subscription.Dispose();
                },
                e =>
                {
                    if (count < 2)
                        o.OnError(e);
                },
                () =>
                {
                    if (count == 1)
                    {
                        o.OnNext(firstValue);
                        o.OnCompleted();
                    }
                    else if (count == 0)
                    {
                        o.OnError(new InvalidOperationException("Sequence contains no elements"));
                    }
                }));
                return subscription;
            });
        }

        public static IObservable<T> WhenSingle<T>(this IObservable<T> source)
        {
            return WhenSingle(source, true);
        }

        public static IObservable<T> WhenSingleOrDefault<T>(this IObservable<T> source)
        {
            return WhenSingle(source, false);
        }

        public static IObservable<Unit> WhenAll<T>(this IObservable<T> source, Func<T, bool> predicate)
        {
            return source.Where(n => !predicate(n)).WhenIsEmpty();
        }

        public static IObservable<Unit> WhenAny<T>(this IObservable<T> source, Func<T, bool> predicate)
        {
            return source.Where(predicate).WhenFirstOrDefault().Select(_ => new Unit());
        }
    }
}
